import tensorflow as tf
from tensorflow import keras
def NN_cae():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(70, 3, padding='SAME', activation=tf.nn.relu),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(40, activation=tf.nn.relu),
        tf.keras.layers.Dense(20, activation=tf.nn.relu),
        tf.keras.layers.Dense(10, activation=tf.nn.relu),
        tf.keras.layers.Dense(1)
    ])
    return model

def loader_cae_nn(middle_out,y_train,batch_size):
    cea_nn_dataset = tf.data.Dataset.from_tensor_slices((middle_out, y_train))
    cea_nn_dataset = cea_nn_dataset.batch(batch_size)
    return cea_nn_dataset
